{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "mnist_ca",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "28S76DVlfCMZ"
      },
      "source": [
        "# [Self-classifying Neural Cellular Automata](https://distill.pub/2020/selforg-mnist/)\n",
        "\n",
        "This notebook contains code to reproduce experiments and figures for the \"Self-classifying Neural Cellular Automata\" article.\n",
        "\n",
        "*Copyright 2020 Google LLC*\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "[https://www.apache.org/licenses/LICENSE-2.0](https://www.apache.org/licenses/LICENSE-2.0)\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "i5wi_r4gyzFr",
        "colab": {}
      },
      "source": [
        "#@title imports and notebook utils\n",
        "%tensorflow_version 2.x\n",
        "\n",
        "import os\n",
        "import io\n",
        "import PIL.Image, PIL.ImageDraw\n",
        "import base64\n",
        "import zipfile\n",
        "import json\n",
        "import requests\n",
        "import numpy as np\n",
        "import matplotlib.pylab as pl\n",
        "import glob\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from IPython.display import Image, HTML, clear_output\n",
        "import tqdm\n",
        "\n",
        "import os\n",
        "os.environ['FFMPEG_BINARY'] = 'ffmpeg'\n",
        "import moviepy.editor as mvp\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\n",
        "clear_output()\n",
        "\n",
        "def np2pil(a):\n",
        "  if a.dtype in [np.float32, np.float64]:\n",
        "    a = np.uint8(np.clip(a, 0, 1)*255)\n",
        "  return PIL.Image.fromarray(a)\n",
        "\n",
        "def imwrite(f, a, fmt=None):\n",
        "  a = np.asarray(a)\n",
        "  if isinstance(f, str):\n",
        "    fmt = f.rsplit('.', 1)[-1].lower()\n",
        "    if fmt == 'jpg':\n",
        "      fmt = 'jpeg'\n",
        "    f = open(f, 'wb')\n",
        "  np2pil(a).save(f, fmt, quality=95)\n",
        "\n",
        "def imencode(a, fmt='jpeg'):\n",
        "  a = np.asarray(a)\n",
        "  if len(a.shape) == 3 and a.shape[-1] == 4:\n",
        "    fmt = 'png'\n",
        "  f = io.BytesIO()\n",
        "  imwrite(f, a, fmt)\n",
        "  return f.getvalue()\n",
        "\n",
        "def im2url(a, fmt='jpeg'):\n",
        "  encoded = imencode(a, fmt)\n",
        "  base64_byte_string = base64.b64encode(encoded).decode('ascii')\n",
        "  return 'data:image/' + fmt.upper() + ';base64,' + base64_byte_string\n",
        "\n",
        "def imshow(a, fmt='jpeg'):\n",
        "  display(Image(data=imencode(a, fmt)))\n",
        "\n",
        "def tile2d(a, w=None):\n",
        "  a = np.asarray(a)\n",
        "  if w is None:\n",
        "    w = int(np.ceil(np.sqrt(len(a))))\n",
        "  th, tw = a.shape[1:3]\n",
        "  pad = (w-len(a))%w\n",
        "  a = np.pad(a, [(0, pad)]+[(0, 0)]*(a.ndim-1), 'constant')\n",
        "  h = len(a)//w\n",
        "  a = a.reshape([h, w]+list(a.shape[1:]))\n",
        "  a = np.rollaxis(a, 2, 1).reshape([th*h, tw*w]+list(a.shape[4:]))\n",
        "  return a\n",
        "\n",
        "def zoom(img, scale=4):\n",
        "  img = np.repeat(img, scale, 0)\n",
        "  img = np.repeat(img, scale, 1)\n",
        "  return img\n",
        "\n",
        "class VideoWriter:\n",
        "  def __init__(self, filename, fps=30.0, **kw):\n",
        "    self.writer = None\n",
        "    self.params = dict(filename=filename, fps=fps, **kw)\n",
        "\n",
        "  def add(self, img):\n",
        "    img = np.asarray(img)\n",
        "    if self.writer is None:\n",
        "      h, w = img.shape[:2]\n",
        "      self.writer = FFMPEG_VideoWriter(size=(w, h), **self.params)\n",
        "    if img.dtype in [np.float32, np.float64]:\n",
        "      img = np.uint8(img.clip(0, 1)*255)\n",
        "    if len(img.shape) == 2:\n",
        "      img = np.repeat(img[..., None], 3, -1)\n",
        "    if len(img.shape) == 3 and img.shape[-1] == 4:\n",
        "      img = img[..., :3] * img[..., 3, None]\n",
        "    self.writer.write_frame(img)\n",
        "\n",
        "  def close(self):\n",
        "    if self.writer:\n",
        "      self.writer.close()\n",
        "\n",
        "  def __enter__(self):\n",
        "    return self\n",
        "\n",
        "  def __exit__(self, *kw):\n",
        "    self.close()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zmOPQeGLt9Uv"
      },
      "source": [
        "##Load MNIST"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "2Wo1AApPulqD",
        "colab": {}
      },
      "source": [
        "# @title Generate train/test set from MNIST.\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "x_train = np.array(x_train / 255.0,).astype(np.float32)\n",
        "x_test = np.array(x_test / 255.0,).astype(np.float32)\n",
        "\n",
        "# @title Data generator\n",
        "color_lookup = tf.constant([\n",
        "            [128, 0, 0],\n",
        "            [230, 25, 75],\n",
        "            [70, 240, 240],\n",
        "            [210, 245, 60],\n",
        "            [250, 190, 190],\n",
        "            [170, 110, 40],\n",
        "            [170, 255, 195],\n",
        "            [165, 163, 159],\n",
        "            [0, 128, 128],\n",
        "            [128, 128, 0],\n",
        "            [0, 0, 0], # This is the default for digits.\n",
        "            [255, 255, 255] # This is the background.\n",
        "            ])\n",
        "\n",
        "backgroundWhite = True\n",
        "def color_labels(x, y_pic, disable_black=False, dtype=tf.uint8):\n",
        "  # works for shapes of x [b, r, c] and [r, c]\n",
        "  black_and_white = tf.fill(list(x.shape) + [2], 0.01)\n",
        "  is_gray = tf.cast(x > 0.1, tf.float32)\n",
        "  is_not_gray = 1. - is_gray\n",
        "\n",
        "  y_pic = y_pic * tf.expand_dims(is_gray, -1) # forcibly cancels everything outside of it.\n",
        "  \n",
        "  # if disable_black, make is_gray super low.\n",
        "  if disable_black:\n",
        "    is_gray *= -1e5\n",
        "    # this ensures that you don't draw white in the digits.\n",
        "    is_not_gray += is_gray\n",
        "\n",
        "  bnw_order = [is_gray, is_not_gray] if backgroundWhite else [is_not_gray, is_gray]\n",
        "  black_and_white *= tf.stack(bnw_order, -1)\n",
        "\n",
        "  rgb = tf.gather(\n",
        "      color_lookup,\n",
        "      tf.argmax(tf.concat([y_pic, black_and_white], -1), -1))\n",
        "  if dtype == tf.uint8:\n",
        "    return tf.cast(rgb, tf.uint8)\n",
        "  else:\n",
        "    return tf.cast(rgb, dtype) / 255.\n",
        "\n",
        "def to_ten_dim_label(x, y):\n",
        "  # x shape is [b, r, c]\n",
        "  # y shape is [b]\n",
        "\n",
        "  # y_res shape is [b, r, c, 10]\n",
        "  y_res = np.zeros(list(x.shape) + [10])\n",
        "  # broadcast y to match x shape:\n",
        "  y_expanded = np.broadcast_to(y, x.T.shape).T\n",
        "  y_res[x >= 0.1, y_expanded[x >= 0.1]] = 1.0\n",
        "  return y_res.astype(np.float32)\n",
        "\n",
        "def find_different_numbers(x_set, y_set, y_set_pic, orientation=\"vertical\"):\n",
        "  result_y = []\n",
        "  result_x = []\n",
        "  for i in range(10):\n",
        "    for x, y, y_pic in zip(x_set, y_set, y_set_pic):\n",
        "      if y == i:\n",
        "        result_y.append(color_labels(x, y_pic))\n",
        "        result_x.append(x)\n",
        "        break\n",
        "  assert len(result_y) == 10\n",
        "\n",
        "  result_y = np.concatenate(result_y, axis=0 if orientation == \"vertical\" else 1)\n",
        "  result_x = np.stack(result_x)\n",
        "\n",
        "  return result_y, result_x\n",
        "\n",
        "\n",
        "print(\"Generating y pics...\")\n",
        "y_train_pic = to_ten_dim_label(x_train, y_train)\n",
        "y_test_pic = to_ten_dim_label(x_test, y_test)\n",
        "\n",
        "\n",
        "numbers_legend, x_legend = find_different_numbers(x_train, y_train, y_train_pic)\n",
        "numbers_legend_horiz, _ = find_different_numbers(x_train, y_train, y_train_pic, \"horizontal\")\n",
        "\n",
        "imshow(zoom(numbers_legend_horiz))\n",
        "print(\"Storing x_legend for use in the demo.\")\n",
        "pl.imshow(x_legend.reshape((-1, 28)))\n",
        "import json\n",
        "samples_str = json.dumps(x_legend.tolist())\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "O-ssztk0fwi5"
      },
      "source": [
        "## Cellular Automata parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "zR6I1JONmWBb",
        "colab": {}
      },
      "source": [
        "#@markdown ### Model configuration\n",
        "#@markdown These options configure the model to be used and train in this \n",
        "#@markdown notebook. Please refer to the article for more information.\n",
        "CHANNEL_N = 19 # Number of CA state channels\n",
        "BATCH_SIZE = 16\n",
        "POOL_SIZE = BATCH_SIZE * 10\n",
        "CELL_FIRE_RATE = 0.5\n",
        "\n",
        "MODEL_TYPE = '3 mutating'  #@param ['1 naive', '2 persistent', '3 mutating']\n",
        "LOSS_TYPE = \"l2\"  #@param ['l2', 'ce']\n",
        "ADD_NOISE = \"True\"  #@param ['True', 'False']\n",
        "\n",
        "USE_PATTERN_POOL, MUTATE_POOL = {\n",
        "    '1 naive': (False, False),\n",
        "    '2 persistent': (True, False),\n",
        "    '3 mutating': (True, True)\n",
        "    }[MODEL_TYPE]\n",
        "ADD_NOISE = ADD_NOISE == 'True'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "lCbPFbI_zosW",
        "colab": {}
      },
      "source": [
        "#@title CA model and utils\n",
        "\n",
        "from tensorflow.keras.layers import Conv2D\n",
        "\n",
        "class CAModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, channel_n=CHANNEL_N, fire_rate=CELL_FIRE_RATE,\n",
        "               add_noise=ADD_NOISE):\n",
        "    # CHANNEL_N does *not* include the greyscale channel.\n",
        "    # but it does include the 10 possible outputs.\n",
        "    super().__init__()\n",
        "    self.channel_n = channel_n\n",
        "    self.fire_rate = fire_rate\n",
        "    self.add_noise = add_noise\n",
        "\n",
        "    self.perceive = tf.keras.Sequential([\n",
        "          Conv2D(80, 3, activation=tf.nn.relu, padding=\"SAME\"),\n",
        "      ])\n",
        "\n",
        "    self.dmodel = tf.keras.Sequential([\n",
        "          Conv2D(80, 1, activation=tf.nn.relu),\n",
        "          Conv2D(self.channel_n, 1, activation=None,\n",
        "                       kernel_initializer=tf.zeros_initializer),\n",
        "    ])\n",
        "\n",
        "    self(tf.zeros([1, 3, 3, channel_n + 1]))  # dummy calls to build the model\n",
        "\n",
        "  @tf.function\n",
        "  def call(self, x, fire_rate=None, manual_noise=None):\n",
        "    gray, state = tf.split(x, [1, self.channel_n], -1)\n",
        "    ds = self.dmodel(self.perceive(x))\n",
        "    if self.add_noise:\n",
        "      if manual_noise is None:\n",
        "        residual_noise = tf.random.normal(tf.shape(ds), 0., 0.02)\n",
        "      else:\n",
        "        residual_noise = manual_noise\n",
        "      ds += residual_noise\n",
        "\n",
        "    if fire_rate is None:\n",
        "      fire_rate = self.fire_rate\n",
        "    update_mask = tf.random.uniform(tf.shape(x[:, :, :, :1])) <= fire_rate\n",
        "    living_mask = gray > 0.1\n",
        "    residual_mask = update_mask & living_mask\n",
        "    ds *= tf.cast(residual_mask, tf.float32)\n",
        "    state += ds\n",
        "    \n",
        "    return tf.concat([gray, state], -1)\n",
        "\n",
        "  @tf.function\n",
        "  def initialize(self, images):\n",
        "    state = tf.zeros([tf.shape(images)[0], 28, 28, self.channel_n])\n",
        "    images = tf.reshape(images, [-1, 28, 28, 1])\n",
        "    return tf.concat([images, state], -1)\n",
        "\n",
        "  @tf.function\n",
        "  def classify(self, x):\n",
        "    # The last 10 layers are the classification predictions, one channel\n",
        "    # per class. Keep in mind there is no \"background\" class,\n",
        "    # and that any loss doesn't propagate to \"dead\" pixels.\n",
        "    return x[:,:,:,-10:]\n",
        "\n",
        "CAModel().perceive.summary()\n",
        "CAModel().dmodel.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xDX5HL7VLd0z"
      },
      "source": [
        "# Training and visualization utils"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "IeWf6HeTe8kM",
        "colab": {}
      },
      "source": [
        "#@title Train utils (SamplePool, model export, visualizations)\n",
        "from google.protobuf.json_format import MessageToDict\n",
        "from tensorflow.python.framework import convert_to_constants\n",
        "\n",
        "class SamplePool:\n",
        "  def __init__(self, *, _parent=None, _parent_idx=None, **slots):\n",
        "    self._parent = _parent\n",
        "    self._parent_idx = _parent_idx\n",
        "    self._slot_names = slots.keys()\n",
        "    self._size = None \n",
        "    for k, v in slots.items():\n",
        "      if self._size is None:\n",
        "        self._size = len(v)\n",
        "      assert self._size == len(v)\n",
        "      setattr(self, k, np.asarray(v))\n",
        "\n",
        "  def sample(self, n):\n",
        "    idx = np.random.choice(self._size, n, False)\n",
        "    batch = {k: getattr(self, k)[idx] for k in self._slot_names}\n",
        "    batch = SamplePool(**batch, _parent=self, _parent_idx=idx)\n",
        "    return batch\n",
        "\n",
        "  def commit(self):\n",
        "    for k in self._slot_names:\n",
        "      getattr(self._parent, k)[self._parent_idx] = getattr(self, k)\n",
        "\n",
        "def export_model(ca, base_fn):\n",
        "  ca.save_weights(base_fn)\n",
        "\n",
        "  cf = ca.call.get_concrete_function(\n",
        "      x=tf.TensorSpec([None, None, None, CHANNEL_N+1]),\n",
        "      fire_rate=tf.constant(0.5),\n",
        "      manual_noise=tf.TensorSpec([None, None, None, CHANNEL_N]))\n",
        "  cf = convert_to_constants.convert_variables_to_constants_v2(cf)\n",
        "  graph_def = cf.graph.as_graph_def()\n",
        "  graph_json = MessageToDict(graph_def)\n",
        "  graph_json['versions'] = dict(producer='1.14', minConsumer='1.14')\n",
        "  model_json = {\n",
        "      'format': 'graph-model',\n",
        "      'modelTopology': graph_json,\n",
        "      'weightsManifest': [],\n",
        "  }\n",
        "  with open(base_fn+'.json', 'w') as f:\n",
        "    json.dump(model_json, f)\n",
        "\n",
        "def classify_and_color(ca, x, disable_black=False):\n",
        "  return color_labels(\n",
        "      x[:,:,:,0], ca.classify(x), disable_black, dtype=tf.float32)\n",
        "\n",
        "\n",
        "def generate_tiled_figures(figures, fade_by=0.1):\n",
        "  tiled_pool = tile2d(figures)\n",
        "  fade_sz = int(tiled_pool.shape[0] * fade_by)\n",
        "  fade = np.linspace(1.0, 0.0, fade_sz)\n",
        "  ones = np.ones(fade_sz) \n",
        "  tiled_pool[:, :fade_sz] += (-tiled_pool[:, :fade_sz] + ones[None, :, None]) * fade[None, :, None] \n",
        "  tiled_pool[:, -fade_sz:] += (-tiled_pool[:, -fade_sz:] + ones[None, :, None]) * fade[None, ::-1, None]\n",
        "  tiled_pool[:fade_sz, :] += (-tiled_pool[:fade_sz, :] + ones[:, None, None]) * fade[:, None, None]\n",
        "  tiled_pool[-fade_sz:, :] += (-tiled_pool[-fade_sz:, :] + ones[:, None, None]) * fade[::-1, None, None]\n",
        "  return tiled_pool\n",
        "\n",
        "def generate_pool_figures(ca, pool, step_i):\n",
        "  tiled_pool = tile2d(classify_and_color(ca, pool.x))\n",
        "  fade = np.linspace(1.0, 0.0, 72)\n",
        "  ones = np.ones(72) \n",
        "  tiled_pool[:, :72] += (-tiled_pool[:, :72] + ones[None, :, None]) * fade[None, :, None] \n",
        "  tiled_pool[:, -72:] += (-tiled_pool[:, -72:] + ones[None, :, None]) * fade[None, ::-1, None]\n",
        "  tiled_pool[:72, :] += (-tiled_pool[:72, :] + ones[:, None, None]) * fade[:, None, None]\n",
        "  tiled_pool[-72:, :] += (-tiled_pool[-72:, :] + ones[:, None, None]) * fade[::-1, None, None]\n",
        "  imwrite('train_log/%04d_pool.jpg'%step_i, tiled_pool)\n",
        "\n",
        "def visualize_batch(ca, x0, x, step_i):\n",
        "  vis0 = np.hstack(classify_and_color(ca, x0).numpy())\n",
        "  vis1 = np.hstack(classify_and_color(ca, x).numpy())\n",
        "  vis = np.vstack([vis0, vis1])\n",
        "  imwrite('train_log/batches_%04d.jpg'%step_i, vis)\n",
        "  print('batch (before/after):')\n",
        "  imshow(vis)\n",
        "\n",
        "def plot_loss(loss_log):\n",
        "  pl.figure(figsize=(10, 4))\n",
        "  pl.title('Loss history (log10)')\n",
        "  pl.plot(np.log10(loss_log), '.', alpha=0.1)\n",
        "  pl.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "9I8YVbzwO-T9",
        "colab": {}
      },
      "source": [
        "# @title Evaluation functions\n",
        "\n",
        "def eval_perform_steps(ca, x, yt, num_steps):\n",
        "  yt_label = tf.argmax(yt, axis=-1)\n",
        "\n",
        "  live_mask = x[..., 0] > 0.1\n",
        "  live_mask_fl = tf.expand_dims(tf.cast(live_mask, tf.float32), -1)\n",
        "  dead_channel = tf.cast(x[..., :1] <= 0.1, tf.float32)\n",
        "\n",
        "  # for now the metric is aggregating everything.\n",
        "  total_count = tf.reduce_sum(tf.cast(live_mask, tf.float32))\n",
        "\n",
        "  avg_accuracy_list = []\n",
        "  avg_total_agreement_list = []\n",
        "  for _ in range(1, num_steps + 1):\n",
        "    x = ca(x)\n",
        "\n",
        "    y = ca.classify(x)\n",
        "    y_label = tf.argmax(y, axis=-1)\n",
        "\n",
        "    correct = tf.equal(y_label,  yt_label) & live_mask\n",
        "    total_correct = tf.reduce_sum(tf.cast(correct, tf.float32))\n",
        "    avg_accuracy_list.append((total_correct/total_count * 100).numpy().item())\n",
        "\n",
        "    # agreement metrics\n",
        "    # Important to exclude dead cells:\n",
        "    y = y * live_mask_fl\n",
        "    y_label_plus_mask = tf.argmax(tf.concat([y, dead_channel], -1), axis=-1)\n",
        "    all_counts = []\n",
        "    for idx in range(10):\n",
        "      count_i = tf.reduce_sum(\n",
        "          tf.cast(tf.equal(y_label_plus_mask, idx), tf.int32), axis=[1,2])\n",
        "      all_counts.append(count_i)\n",
        "    all_counts_t = tf.stack(all_counts, 1)\n",
        "    # Now the trick is that if there is a total agreement, their sum is the same\n",
        "    # as their max.\n",
        "    equality = tf.equal(tf.reduce_max(all_counts_t, axis=1),\n",
        "                        tf.reduce_sum(all_counts_t, axis=1))\n",
        "    sum_agreement = tf.reduce_sum(tf.cast(equality, tf.float32))\n",
        "    avg_total_agreement_list.append(sum_agreement.numpy().item() / y.shape[0] * 100)\n",
        "\n",
        "  return avg_accuracy_list, avg_total_agreement_list\n",
        "\n",
        "def eval_batch_fn(ca, x_test, y_test_pic, num_steps, mutate):\n",
        "  x = ca.initialize(x_test)\n",
        "  yt = y_test_pic\n",
        "\n",
        "  avg_acc_l_1, avg_tot_agr_l_1 = eval_perform_steps(ca, x, yt, num_steps)\n",
        "  if not mutate:\n",
        "    return avg_acc_l_1, avg_tot_agr_l_1\n",
        "  # Accuracy after mutation!\n",
        "  new_idx = np.random.randint(0, x_test.shape[0]-1, size=x_test.shape[0])\n",
        "  new_x, yt = x_test[new_idx], y_test_pic[new_idx]\n",
        "  new_x = tf.reshape(new_x, [-1, 28, 28, 1])\n",
        "  mutate_mask = tf.cast(new_x > 0.1, tf.float32)\n",
        "\n",
        "  x = tf.concat([new_x, x[:,:,:,1:] * mutate_mask], -1)\n",
        "\n",
        "  avg_acc_l_2, avg_tot_agr_l_2 = eval_perform_steps(ca, x, yt, num_steps)\n",
        "\n",
        "  return avg_acc_l_1 + avg_acc_l_2, avg_tot_agr_l_1 + avg_tot_agr_l_2\n",
        "\n",
        "def eval_all(ca, num_steps, mutate):\n",
        "  all_accuracies = []\n",
        "  all_agreements = []\n",
        "\n",
        "  # total test set is 10000\n",
        "  num_batches = 10\n",
        "  eval_bs = 10000 // num_batches\n",
        "  for i in tqdm.trange(num_batches):\n",
        "    x_set = x_test[eval_bs*i:eval_bs*(i+1)]\n",
        "    y_set = y_test_pic[eval_bs*i:eval_bs*(i+1)]\n",
        "    acc_i, agr_i = eval_batch_fn(ca, x_set, y_set, num_steps, mutate)\n",
        "    all_accuracies.append(acc_i)\n",
        "    all_agreements.append(agr_i)\n",
        "\n",
        "  all_accuracies = [sum(l)/num_batches for l in zip(*all_accuracies)]\n",
        "  all_agreements = [sum(l)/num_batches for l in zip(*all_agreements)]\n",
        "  return all_accuracies, all_agreements\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qdfCyWWwZFWq"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "ak5rBmbxmHV7",
        "colab": {}
      },
      "source": [
        "# Initialize things for a new training run\n",
        "ca = CAModel()\n",
        "\n",
        "def individual_l2_loss(x, y):\n",
        "  t = y - ca.classify(x)\n",
        "  return tf.reduce_sum(t**2, [1, 2, 3]) / 2\n",
        "\n",
        "def batch_l2_loss(x, y):\n",
        "  return tf.reduce_mean(individual_l2_loss(x, y))\n",
        "\n",
        "def batch_ce_loss(x, y):\n",
        "  one_hot = tf.argmax(y, axis=-1)\n",
        "  # It's ok even if the loss is computed on \"dead\" cells. Anyway they shouldn't\n",
        "  # get any gradient propagated through there.\n",
        "  return tf.compat.v1.losses.sparse_softmax_cross_entropy(one_hot, x)\n",
        "\n",
        "assert LOSS_TYPE in [\"l2\", \"ce\"]\n",
        "loss_fn = batch_l2_loss if LOSS_TYPE == \"l2\" else batch_ce_loss\n",
        "\n",
        "loss_log = []\n",
        "\n",
        "lr = 1e-3\n",
        "lr_sched = tf.keras.optimizers.schedules.PiecewiseConstantDecay(\n",
        "      [30000, 70000], [lr, lr*0.1, lr*0.01])\n",
        "trainer = tf.keras.optimizers.Adam(lr_sched)\n",
        "\n",
        "starting_indexes = np.random.randint(0, x_train.shape[0]-1, size=POOL_SIZE)\n",
        "pool = SamplePool(x=ca.initialize(x_train[starting_indexes]).numpy(), y=y_train_pic[starting_indexes])\n",
        "\n",
        "!mkdir -p train_log && rm -f train_log/*"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "QzP_vDchq0d9",
        "colab": {}
      },
      "source": [
        "# Training loop\n",
        "\n",
        "@tf.function\n",
        "def train_step(x, y):\n",
        "  iter_n = 20\n",
        "  with tf.GradientTape() as g:\n",
        "    for i in tf.range(iter_n):\n",
        "      x = ca(x)\n",
        "    loss = batch_l2_loss(x, y)\n",
        "  grads = g.gradient(loss, ca.weights)\n",
        "  grads = [g/(tf.norm(g)+1e-8) for g in grads]\n",
        "  trainer.apply_gradients(zip(grads, ca.weights))\n",
        "  return x, loss\n",
        "\n",
        "for i in range(1, 100000+1):\n",
        "  if USE_PATTERN_POOL:\n",
        "    batch = pool.sample(BATCH_SIZE)\n",
        "    x0 = np.copy(batch.x)\n",
        "    y0 = batch.y\n",
        "    # we want half of them new. We remove 1/4 from the top and 1/4 from the\n",
        "    # bottom.\n",
        "    q_bs = BATCH_SIZE // 4\n",
        "\n",
        "    new_idx = np.random.randint(0, x_train.shape[0]-1, size=q_bs)\n",
        "    x0[:q_bs] = ca.initialize(x_train[new_idx])\n",
        "    y0[:q_bs] = y_train_pic[new_idx]\n",
        "\n",
        "    new_idx = np.random.randint(0, x_train.shape[0]-1, size=q_bs)\n",
        "    new_x, new_y = x_train[new_idx], y_train_pic[new_idx]\n",
        "    if MUTATE_POOL:\n",
        "      new_x = tf.reshape(new_x, [q_bs, 28, 28, 1])\n",
        "      mutate_mask = tf.cast(new_x > 0.1, tf.float32)\n",
        "      mutated_x = tf.concat([new_x, x0[-q_bs:,:,:,1:] * mutate_mask], -1)\n",
        "\n",
        "      x0[-q_bs:] = mutated_x\n",
        "      y0[-q_bs:] = new_y\n",
        "\n",
        "    else:\n",
        "      x0[-q_bs:] = ca.initialize(new_x)\n",
        "      y0[-q_bs:] = new_y\n",
        "\n",
        "  else:\n",
        "    b_idx = np.random.randint(0, x_train.shape[0]-1, size=BATCH_SIZE)\n",
        "    x0 = ca.initialize(x_train[b_idx])\n",
        "    y0 = y_train_pic[b_idx]\n",
        "\n",
        "  x, loss = train_step(x0, y0)\n",
        "\n",
        "  if USE_PATTERN_POOL:\n",
        "    batch.x[:] = x\n",
        "    batch.y[:] = y0 # This gets reordered, so you need to change it.\n",
        "    batch.commit()\n",
        "\n",
        "  step_i = len(loss_log)\n",
        "  loss_log.append(loss.numpy())\n",
        "\n",
        "  if step_i%100 == 0:\n",
        "    generate_pool_figures(ca, pool, step_i)\n",
        "  if step_i%200 == 0:\n",
        "    clear_output()\n",
        "    visualize_batch(ca, x0, x, step_i)\n",
        "    plot_loss(loss_log)\n",
        "  if step_i%10000 == 0:\n",
        "    export_model(ca, 'train_log/%07d'%step_i)\n",
        "\n",
        "  print('\\r step: %d, log10(loss): %.3f'%(len(loss_log), np.log10(loss)), end='')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "2vhKajnuxINp",
        "colab": {}
      },
      "source": [
        "# useful code if you end up interrupting the run.\n",
        "print(step_i)\n",
        "export_model(ca, 'train_log/%07d'%step_i)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "nqvkfl9W4ODI",
        "colab": {}
      },
      "source": [
        "# @title eval metrics\n",
        "\n",
        "eval_batch = 1000\n",
        "num_iters = 10\n",
        "\n",
        "all_accuracies, all_agreements = eval_all(ca, num_steps=200, mutate=True)\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.title('Average cell accuracy over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average cell accuracy (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_accuracies, label=\"ca\")\n",
        "pl.legend()\n",
        "pl.show()\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.title('Average total agreement across batch over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average total agreement (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_agreements, label=\"ca\")\n",
        "pl.legend()\n",
        "pl.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "QeXZKb5v2gxj",
        "colab": {}
      },
      "source": [
        "#@title training progress (batches)\n",
        "frames = sorted(glob.glob('train_log/batches_*.jpg'))\n",
        "mvp.ImageSequenceClip(frames, fps=10.0).write_videofile('batches.mp4')\n",
        "mvp.ipython_display('batches.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "B4JAbAJf6Alw",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "#@title pool contents\n",
        "#!rm -f train_log/*_pool.jpg\n",
        "\n",
        "frames = sorted(glob.glob('train_log/*_pool.jpg'))[:80]\n",
        "mvp.ImageSequenceClip(frames, fps=5.0).write_videofile('pool.mp4')\n",
        "mvp.ipython_display('pool.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "u57S5tCxhcNp"
      },
      "source": [
        "## pretrained models\n",
        "\n",
        "Please run the cell below to download pretrained models that are used to generate the subsequent figures. They are also used in the demo section."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MqxOYKrFS8fU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget -O models.zip 'https://github.com/google-research/self-organising-systems/blob/master/assets/mnist_ca/models.zip?raw=true'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ONAhIfBh5X0N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip -oq \"models.zip\" -d \"models\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "wiGl7S0E6-OA",
        "colab": {}
      },
      "source": [
        "def get_exp_path(\n",
        "    prefix, use_sample_pool, mutate_pool, loss_type, add_noise):\n",
        "  path = prefix\n",
        "  path += 'use_sample_pool_%r mutate_pool_%r '%(use_sample_pool, mutate_pool)\n",
        "  path += 'loss_type_%s '%(loss_type)\n",
        "  path += 'add_noise_%r'%(add_noise)\n",
        "  path += '/0100000'\n",
        "  return path\n",
        "\n",
        "def get_model(use_sample_pool=True, mutate_pool=True, loss_type=\"l2\", add_noise=True,\n",
        "              prefix=\"models/\", output='model'):\n",
        "  path = get_exp_path(\n",
        "      prefix, use_sample_pool, mutate_pool, loss_type, add_noise)\n",
        "  assert output in ['model', 'json']\n",
        "  if output == 'model':\n",
        "    ca = CAModel(add_noise=add_noise)\n",
        "    ca.load_weights(path)\n",
        "    return ca\n",
        "  elif output == 'json':\n",
        "    return open(path+'.json', 'r').read()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QAscSKkRaFwp"
      },
      "source": [
        "# Figures"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "HgkWl2KD8f-A",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title mosaic of pictures\n",
        "\n",
        "num_digits = 144\n",
        "\n",
        "indexes = np.random.randint(0, x_train.shape[0]-1, size=num_digits)\n",
        "figures = color_labels(x_train[indexes], y_train_pic[indexes], dtype=tf.float32)\n",
        "tile_digits = generate_tiled_figures(figures, fade_by=0.1)\n",
        "\n",
        "imshow(zoom(tile_digits))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Itv0WEuRVWoh"
      },
      "source": [
        "## visualize runs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qdpBqDXySMUn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget -O mnist_slider.png 'https://github.com/google-research/self-organising-systems/blob/master/assets/mnist_ca/mnist_slider.png?raw=true'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "co32HAdaVZWt",
        "colab": {}
      },
      "source": [
        "def make_run_videos(ca, num_steps, eval_bs, prefix, disable_black=False):\n",
        "  new_idx = np.random.randint(0, x_train.shape[0]-1, size=eval_bs)\n",
        "  x = ca.initialize(x_train[new_idx])\n",
        "  frames = []\n",
        "  with VideoWriter(prefix + \".mp4\") as vid:\n",
        "    slider = PIL.Image.open(\"mnist_slider.png\")\n",
        "    for i in tqdm.trange(-1, num_steps):\n",
        "      if i == -1:\n",
        "        image = zoom(tile2d(classify_and_color(ca, x, disable_black=False)), scale=2)\n",
        "      else:\n",
        "        if i == num_steps//2:\n",
        "          # then mutate\n",
        "          new_idx = np.random.randint(0, x_train.shape[0]-1, size=eval_bs)\n",
        "          new_x = x_train[new_idx]\n",
        "          new_x = tf.reshape(new_x, [eval_bs, 28, 28, 1])\n",
        "          mutate_mask = tf.cast(new_x > 0.1, tf.float32)\n",
        "          x = tf.concat([new_x, x[:,:,:,1:] * mutate_mask], -1)\n",
        "        x = ca(x)\n",
        "        image = zoom(tile2d(classify_and_color(ca, x, disable_black=disable_black)), scale=2)\n",
        "      vis_extended = np.concatenate((image, np.ones((86, image.shape[1], 3))), axis=0) \n",
        "      im = np.uint8(vis_extended*255)\n",
        "      im = PIL.Image.fromarray(im)\n",
        "      im.paste(slider, box=(0, image.shape[0]+20))\n",
        "      draw = PIL.ImageDraw.Draw(im)\n",
        "      p_x = 3+(((image.shape[1]-5-3)/num_steps)*i)\n",
        "      draw.rectangle([p_x, image.shape[0]+21, p_x+5, image.shape[0]+42], fill=\"#434343bd\")\n",
        "      vid.add(np.uint8(im))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "FgKYIUxjGNBy",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title visualize CE runs\n",
        "eval_bs= 100\n",
        "num_steps = 400 \n",
        "\n",
        "ca = get_model(add_noise=False, loss_type='ce')\n",
        "make_run_videos(ca, num_steps, eval_bs, \"ce_runs\", disable_black=True)\n",
        "mvp.ipython_display('ce_runs.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "P35Un4ZxYFIq",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title visualize L2 runs\n",
        "eval_bs= 100\n",
        "num_steps = 400\n",
        "\n",
        "ca = get_model(add_noise=True, loss_type='l2')\n",
        "make_run_videos(ca, num_steps, eval_bs, \"l2_runs\")\n",
        "mvp.ipython_display('l2_runs.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "np2Z1XzeC-Je",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title CE eval for experiment 1\n",
        "# @title comparing\n",
        "\n",
        "ca = get_model(add_noise=False, loss_type='ce')\n",
        "label_1 = 'Cross Entropy'\n",
        "\n",
        "all_accuracies, all_agreements = eval_all(ca, num_steps=200, mutate=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "U3PRdvs5NyN0",
        "colab": {}
      },
      "source": [
        "pl.figure(figsize=(10, 4))\n",
        "\n",
        "pl.title('Average cell accuracy over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average cell accuracy (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_accuracies, label=label_1)\n",
        "pl.legend()\n",
        "pl.show()\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.title('Average total agreement across batch over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average total agreement (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_agreements, label=label_1)\n",
        "pl.legend()\n",
        "pl.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "0pg8b2FIQAMm",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title comparing CE, L2, L2+noise\n",
        "print(\"Evaluating CE:\", flush=True)\n",
        "ca = get_model(add_noise=False, loss_type='ce')\n",
        "all_accuracies_1, all_agreements_1 = eval_all(ca, num_steps=200, mutate=True)\n",
        "\n",
        "print(\"Evaluating L2:\", flush=True)\n",
        "ca = get_model(add_noise=False, loss_type='l2')\n",
        "all_accuracies_2, all_agreements_2 = eval_all(ca, num_steps=200, mutate=True)\n",
        "\n",
        "print(\"Evaluating L2 with noise:\", flush=True)\n",
        "ca = get_model(add_noise=True, loss_type='l2')\n",
        "all_accuracies_3, all_agreements_3 = eval_all(ca, num_steps=200, mutate=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "e5JvuWuU-N87",
        "colab": {}
      },
      "source": [
        "def get_highest_point(arr):\n",
        "  highest_i = 0\n",
        "  highest_val = 0\n",
        "  for i in range(len(arr)):\n",
        "    if arr[i] > highest_val:\n",
        "      highest_i = i\n",
        "      highest_val = arr[i]\n",
        "  return highest_i, highest_val\n",
        "\n",
        "highest_i, highest_val = get_highest_point(all_accuracies_1[:200])\n",
        "print(\"CE, top accuracy:\", highest_i, highest_val)\n",
        "highest_i, highest_val = get_highest_point(all_agreements_1[:200])\n",
        "print(\"CE, top agreement:\", highest_i, highest_val)\n",
        "\n",
        "highest_i, highest_val = get_highest_point(all_accuracies_2[:200])\n",
        "print(\"L2, top accuracy:\", highest_i, highest_val)\n",
        "highest_i, highest_val = get_highest_point(all_agreements_2[:200])\n",
        "print(\"L2, top agreement:\", highest_i, highest_val)\n",
        "\n",
        "\n",
        "highest_i, highest_val = get_highest_point(all_accuracies_3[:200])\n",
        "print(\"L2 with noise, top accuracy:\", highest_i, highest_val)\n",
        "highest_i, highest_val = get_highest_point(all_agreements_3[:200])\n",
        "print(\"L2 with noise, top agreement:\", highest_i, highest_val)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "G0aMTsxRQ4w-",
        "colab": {}
      },
      "source": [
        "label_1 = 'Cross Entropy'\n",
        "label_2 = 'L2'\n",
        "label_3 = 'L2 + Residual Noise'\n",
        "\n",
        "pl.figure(figsize=(15, 10))\n",
        "pl.subplot(211)\n",
        "pl.title('Average cell accuracy over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average cell accuracy (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_accuracies_1, label=label_1)\n",
        "pl.plot(all_accuracies_2, label=label_2)\n",
        "pl.plot(all_accuracies_3, label=label_3)\n",
        "pl.legend()\n",
        "\n",
        "pl.subplot(212)\n",
        "pl.title('Average total agreement across batch over steps (%)')\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average total agreement (%)')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 100)\n",
        "pl.plot(all_agreements_1, label=label_1)\n",
        "pl.plot(all_agreements_2, label=label_2)\n",
        "pl.plot(all_agreements_3, label=label_3)\n",
        "pl.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "23ezMa38gEE5"
      },
      "source": [
        "## Internal states\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "ok94kFMsGacj",
        "colab": {}
      },
      "source": [
        "# @title Eval magnitude function\n",
        "\n",
        "def eval_mag_perform_steps(ca, x, num_steps):\n",
        "  m = x[...,:1]>0.1\n",
        "  m_fl = tf.cast(m, tf.float32)\n",
        "  m_sum = tf.reduce_sum(m_fl) * 19 # because there are 19 mutable states.\n",
        "\n",
        "  # Update mask is known to be 50%, so we just know to divide the denominator\n",
        "  # by 2:\n",
        "  m_sum_active = m_sum / 2.\n",
        "\n",
        "  state_magnitudes = []\n",
        "  state_delta_magnitudes = []\n",
        "\n",
        "  for _ in range(1, num_steps + 1):\n",
        "    x_prev = x[..., 1:] # ignore the grey channel\n",
        "    x = ca(x)\n",
        "\n",
        "    # State magnitude\n",
        "    # Ideally, these would have to be multiplied by the active mask, but\n",
        "    # we know everything else is zero anyway.\n",
        "    mutable_states = x[..., 1:]\n",
        "    to_plot = tf.reduce_sum(tf.abs(mutable_states)) / m_sum\n",
        "    state_magnitudes.append(to_plot.numpy().item())\n",
        "\n",
        "    # deltas\n",
        "    to_plot = tf.reduce_sum(tf.abs( mutable_states - x_prev)) / m_sum_active\n",
        "    state_delta_magnitudes.append(to_plot.numpy().item())\n",
        "\n",
        "  return state_magnitudes, state_delta_magnitudes\n",
        "\n",
        "def eval_mag_batch_fn(ca, x_test, num_steps, mutate):\n",
        "  x = ca.initialize(x_test)\n",
        "\n",
        "  st_mag_l_1, st_delta_mag_l_1 = eval_mag_perform_steps(ca, x, num_steps)\n",
        "  if not mutate:\n",
        "    return st_mag_l_1, st_delta_mag_l_1\n",
        "  # after mutation!\n",
        "  new_idx = np.random.randint(0, x_test.shape[0]-1, size=x_test.shape[0])\n",
        "  new_x = x_test[new_idx]\n",
        "  new_x = tf.reshape(new_x, [-1, 28, 28, 1])\n",
        "  mutate_mask = tf.cast(new_x > 0.1, tf.float32)\n",
        "\n",
        "  x = tf.concat([new_x, x[:,:,:,1:] * mutate_mask], -1)\n",
        "\n",
        "  st_mag_l_2, st_delta_mag_l_2 = eval_mag_perform_steps(ca, x, num_steps)\n",
        "\n",
        "  return st_mag_l_1 +st_mag_l_2, st_delta_mag_l_1 + st_delta_mag_l_2\n",
        "\n",
        "def eval_magnitude_all(ca, num_steps, mutate):\n",
        "  all_magnitudes = []\n",
        "  all_delta_magnitudes = []\n",
        "\n",
        "  # total test set is 10000\n",
        "  num_batches = 10\n",
        "  eval_bs = 10000 // num_batches\n",
        "  for i in tqdm.trange(num_batches):\n",
        "    x_set = x_test[eval_bs*i:eval_bs*(i+1)]\n",
        "    mag_i, del_i = eval_mag_batch_fn(ca, x_set, num_steps, mutate)\n",
        "    all_magnitudes.append(mag_i)\n",
        "    all_delta_magnitudes.append(del_i)\n",
        "\n",
        "  all_magnitudes = [sum(l)/num_batches for l in zip(*all_magnitudes)]\n",
        "  all_delta_magnitudes = [sum(l)/num_batches for l in zip(*all_delta_magnitudes)]\n",
        "  return all_magnitudes, all_delta_magnitudes\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "GSS20kb7HPJd",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title CE magnitude eval for experiment 1\n",
        "# @title comparing\n",
        "\n",
        "ca = get_model(add_noise=False, loss_type='ce')\n",
        "ce_all_magnitudes, ce_all_delta_magnitudes = eval_magnitude_all(ca, num_steps=200, mutate=True)\n",
        "pl_colorcycle = pl.rcParams['axes.prop_cycle'].by_key()['color']\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average magnitude')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(0, 30)\n",
        "pl.plot(ce_all_magnitudes, label=\"CE: state\", color=pl_colorcycle[0])\n",
        "pl.plot(ce_all_delta_magnitudes, label=\"CE: residual\", color=pl_colorcycle[0], \n",
        "        linestyle=\"--\")\n",
        "\n",
        "pl.legend()\n",
        "pl.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "Q5GSa39kH-Wo",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title L2 magnitude eval for experiment 2\n",
        "# @title comparing\n",
        "\n",
        "ca = get_model(add_noise=False, loss_type='l2')\n",
        "\n",
        "l2_all_magnitudes, l2_all_delta_magnitudes = eval_magnitude_all(ca, num_steps=200, mutate=True)\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.xlabel('Number of steps')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(1e-2, 4e1)\n",
        "pl.yscale(\"log\")\n",
        "pl.plot(l2_all_magnitudes, label=\"Average state magnitude\")\n",
        "pl.plot(l2_all_delta_magnitudes, label=\"Average residual magnitude\")\n",
        "\n",
        "pl.legend()\n",
        "pl.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "BY7F_HVUIBOt",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title L2+ noise magnitude eval for experiment 2\n",
        "\n",
        "ca = get_model(add_noise=True, loss_type='l2')\n",
        "\n",
        "l2n_all_magnitudes, l2n_all_delta_magnitudes = eval_magnitude_all(ca, num_steps=200, mutate=True)\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.xlabel('Number of steps')\n",
        "pl.xlim(0, 400)\n",
        "pl.plot(l2n_all_magnitudes, label=\"Average state magnitude\")\n",
        "pl.plot(l2n_all_delta_magnitudes, label=\"Average residual magnitude\")\n",
        "\n",
        "pl.legend()\n",
        "pl.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "b_zLgXamIFlB",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title Observed state magnitudes of CE, L2, and L2 + noise\n",
        "\n",
        "pl_colorcycle = pl.rcParams['axes.prop_cycle'].by_key()['color']\n",
        "\n",
        "pl.figure(figsize=(10, 4))\n",
        "pl.xlabel('Number of steps')\n",
        "pl.ylabel('Average magnitude')\n",
        "pl.xlim(0, 400)\n",
        "pl.ylim(1e-2, 4e1)\n",
        "pl.yscale(\"log\")\n",
        "pl.plot(ce_all_magnitudes, label=\"CE: state\", color=pl_colorcycle[0])\n",
        "pl.plot(ce_all_delta_magnitudes, label=\"CE: residual\", color=pl_colorcycle[0], \n",
        "        linestyle=\"--\")\n",
        "\n",
        "pl.plot(l2_all_magnitudes, label=\"L2: state\", color=pl_colorcycle[1])\n",
        "pl.plot(l2_all_delta_magnitudes, label=\"L2: residual\", color=pl_colorcycle[1],\n",
        "        linestyle=\"--\")\n",
        "\n",
        "pl.plot(l2n_all_magnitudes, label=\"L2 + Noise: state\", color=pl_colorcycle[2])\n",
        "pl.plot(l2n_all_delta_magnitudes, label=\"L2 + Noise: residual\", color=pl_colorcycle[2],\n",
        "        linestyle=\"--\")\n",
        "\n",
        "pl.legend(loc=\"lower right\")\n",
        "pl.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JZnHjb4kwOFQ"
      },
      "source": [
        "## Internal states video for L2 model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "328jOjbUIuVe",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title internal states utils \n",
        "from skimage.transform import resize\n",
        "\n",
        "def gen_clip_squasher(max_magnitude):\n",
        "  def squasher(vis):\n",
        "    vis = np.clip(vis, -max_magnitude, max_magnitude)\n",
        "    return vis / max_magnitude\n",
        "  return squasher\n",
        "\n",
        "def vis_internal_states(x, squasher):\n",
        "  vis = x.numpy().transpose([2, 0, 1])\n",
        "  vis = tile2d(vis, 10)\n",
        "  vis = squasher(vis)\n",
        "  vis = zoom(pl.cm.RdBu_r((vis / (np.pi / 2)) / 2. + 0.5), 2)\n",
        "  return vis\n",
        "\n",
        "def get_color_bar(unsquasher, num_ticks, pixel_width=500):\n",
        "  values = np.linspace(0, 1, 256)\n",
        "  values2D = np.vstack((values, values))\n",
        "  fig, ax = pl.subplots()\n",
        "  fig.set_dpi(100)\n",
        "  fig.set_figheight((pixel_width/100) * 0.05)\n",
        "  fig.set_figwidth(pixel_width/100)\n",
        "  ax.imshow(values2D, aspect='auto', cmap=pl.cm.get_cmap(\"RdBu_r\"))\n",
        "  ax.set_xticks(np.linspace(0,255,num_ticks))\n",
        "  ax.get_yaxis().set_visible(False)\n",
        "  ax.xaxis.tick_top()\n",
        "  ax.spines['right'].set_visible(False)\n",
        "  ax.spines['bottom'].set_visible(False)\n",
        "  ax.spines['left'].set_visible(False)\n",
        "  ax.spines['top'].set_visible(False)\n",
        "  labels = np.linspace(0,1,num_ticks)\n",
        "  labels = (np.pi / 2) * (2.0 * (labels - 0.5)) \n",
        "  labels = unsquasher(labels)\n",
        "  labels = [\"%.1f\" % x for x in labels]\n",
        "  labels[0] = \"-∞\"\n",
        "  labels[-1] = \"∞\"\n",
        "  ax.set_xticklabels(labels)\n",
        "  fig.savefig(\"colorbar.jpg\", bbox_inches='tight')\n",
        "  pl.close()\n",
        "  plt_image = PIL.Image.open(\"colorbar.jpg\")\n",
        "  final_output = PIL.Image.new(\"RGB\", (pixel_width, plt_image.height), (255, 255, 255))\n",
        "  final_output.paste(plt_image, (112+int(((pixel_width-112)-plt_image.width)//2), 0))\n",
        "  return final_output \n",
        "\n",
        "def vis_for_video(x, squasher, legend=None):\n",
        "  vis = vis_internal_states(x, squasher)\n",
        "  current_prediction = classify_and_color(ca, tf.expand_dims(x, 0), disable_black=True)[0]\n",
        "  above_part = zoom(current_prediction, 10)\n",
        "  above_part = resize(above_part, (vis.shape[1], vis.shape[1]))\n",
        "  # add A channel\n",
        "  above_part = np.concatenate([above_part, np.ones_like(above_part[...,:1])], -1)\n",
        "  vis = np.vstack([above_part, vis])\n",
        "  if legend is not None:\n",
        "    l_h = legend.shape[1] * vis.shape[0] // legend.shape[0]\n",
        "    if l_h % 2 == 1:  l_h += 1\n",
        "    legend = resize(legend, (vis.shape[0], l_h))\n",
        "    vis = np.hstack([vis, legend])\n",
        "  return vis\n",
        "\n",
        "def vis_for_video_horiz(x, squasher):\n",
        "  vis = vis_internal_states(x, squasher)\n",
        "  current_prediction = classify_and_color(ca, tf.expand_dims(x, 0), disable_black=True)[0]\n",
        "  left_part = zoom(current_prediction, 4)\n",
        "  # add A channel\n",
        "  left_part = np.concatenate([left_part, np.ones_like(left_part[...,:1])], -1)\n",
        "  vis = np.hstack([left_part, vis])\n",
        "  return vis"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R6D8oYv6XCCd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget -O states_slider.png 'https://github.com/google-research/self-organising-systems/blob/master/assets/mnist_ca/states_slider.png?raw=true'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "7lcOXNfF4K-c",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "# @title L2 internal states\n",
        "figure_name = \"l2n_horiz_states\"\n",
        "\n",
        "legend = numbers_legend\n",
        "# need to add the A channel.\n",
        "legend = np.concatenate([legend, np.ones_like(legend[...,:1])], -1)\n",
        "\n",
        "#squasher = gen_clip_squasher(2.0)\n",
        "squasher = np.arctan\n",
        "unsquasher = np.tan\n",
        "\n",
        "placement_style = \"horiz\" # \"vertic\"\n",
        "\n",
        "num_mutations = 5\n",
        "speedup_after = 50\n",
        "speedup_by = 10\n",
        "\n",
        "x_set = x_test\n",
        "y_set = y_test_pic\n",
        "\n",
        "ca = get_model(add_noise=True, loss_type='l2')\n",
        "frames = []\n",
        "width = 2*(28*10 + 2*28)\n",
        "\n",
        "colorbar = get_color_bar(unsquasher, 17, width)\n",
        "states_slider = PIL.Image.open(\"states_slider.png\")\n",
        "\n",
        "progress = 0\n",
        "with VideoWriter(figure_name + \".mp4\", fps=20.0) as vid:\n",
        "  for m in tqdm.trange(num_mutations + 1):\n",
        "    if m == 0:\n",
        "      b_idx = np.random.randint(0, x_set.shape[0]-1, size=1)\n",
        "      x0 = ca.initialize(x_train[b_idx])\n",
        "    else:\n",
        "      # mutate\n",
        "      b_idx = np.random.randint(0, x_set.shape[0]-1, size=1)\n",
        "      new_x = x_set[b_idx]\n",
        "      new_x = tf.reshape(new_x, [-1, 28, 28, 1])\n",
        "      mutate_mask = tf.cast(new_x > 0.1, tf.float32)\n",
        "      x0 = tf.concat([new_x, x0[:,:,:,1:] * mutate_mask], -1)\n",
        "    \n",
        "    for i in range(0, 200):\n",
        "      if i < speedup_after or i % speedup_by == 0:\n",
        "        vis = (vis_for_video_horiz(x0[0], squasher) if placement_style == \"horiz\" else\n",
        "          vis_for_video(x0[0], squasher, legend))\n",
        "        states_slider_progress = states_slider.copy()\n",
        "        draw = PIL.ImageDraw.Draw(states_slider_progress)\n",
        "        p_x = int(progress/((num_mutations+1) * 200) * width)\n",
        "        draw.rectangle([p_x, 0, p_x+5, states_slider_progress.height], fill=\"#434343bd\")\n",
        "        vid.add(np.vstack((np.asarray(colorbar), \n",
        "                           np.uint8(vis[..., :3]*255.0),\n",
        "                           np.uint8(np.ones((10, width, 3))*255.0), \n",
        "                           np.asarray(states_slider_progress))))\n",
        "      progress += 1\n",
        "      x0 = ca(x0)\n",
        "\n",
        "mvp.ipython_display(figure_name + '.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "9M-oDuhea7bR"
      },
      "source": [
        "# Demo\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "HkeuWM7VVSGb",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Create the data for the distill demo\n",
        "\n",
        "def create_numbers_legend(x_set, y_set, y_set_pic):\n",
        "  result_y = []\n",
        "  for i in range(10):\n",
        "    for x, y, y_pic in zip(x_set, y_set, y_set_pic):\n",
        "      if y == i:\n",
        "        result_y.append(color_labels(x, y_pic))\n",
        "        break\n",
        "  assert len(result_y) == 10\n",
        "  return tile2d(np.stack(result_y), 5)\n",
        "\n",
        "\n",
        "def find_several_numbers(x_set, y_set, num_each):\n",
        "  result = []\n",
        "  for i in range(10):\n",
        "    i_numbers = []\n",
        "    n = 0\n",
        "    for x, y in zip(x_set, y_set):\n",
        "      if y == i:\n",
        "        i_numbers.append(x)\n",
        "        n += 1\n",
        "        if n == num_each: break\n",
        "    result.append(np.stack(i_numbers))\n",
        "  return np.stack(result)\n",
        "\n",
        "demo_numbers_legend = create_numbers_legend(x_train, y_train, y_train_pic)\n",
        "\n",
        "n = 20\n",
        "demo_samples = find_several_numbers(x_train, y_train, n)\n",
        "print(demo_samples.shape)\n",
        "demo_samples = np.concatenate([tile2d(dsi, n) for dsi in demo_samples], 0)\n",
        "print(demo_samples.shape)\n",
        "\n",
        "imshow(zoom(demo_numbers_legend))\n",
        "im = PIL.Image.fromarray(zoom(demo_numbers_legend))\n",
        "im.save(\"demo_numbers_legend.png\")\n",
        "\n",
        "print(\"Storing demo_samples for use in the demo.\")\n",
        "import json\n",
        "many_samples_str = json.dumps(demo_samples.tolist())\n",
        "print(len(many_samples_str))\n",
        "\n",
        "pl.imshow(demo_samples)\n",
        "demo_samples *= 255.\n",
        "im = PIL.Image.fromarray(demo_samples.astype(np.uint8)).convert(\"L\")\n",
        "im.save(\"demo_samples.png\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "4C_0Iusaw2TD",
        "colab": {}
      },
      "source": [
        "#@title TensorFlow.js Demo small {run:\"auto\", vertical-output: true}\n",
        "#@markdown Select \"CHECKPOINT\" model to load the checkpoint created by running \n",
        "#@markdown cells from the \"Training\" section of this notebook.\n",
        "#@markdown Technical note: CE models should be rendered differently to avoid\n",
        "#@markdown black pixels showing for low magnitude.\n",
        "import IPython.display\n",
        "\n",
        "model_source = \"LOAD\"  #@param ['CHECKPOINT', 'LOAD']\n",
        "model_type = '3 mutating'  #@param ['1 naive', '2 persistent', '3 mutating']\n",
        "loss_type = \"l2\"  #@param ['l2', 'ce']\n",
        "add_noise = \"True\"  #@param ['True', 'False']\n",
        "#quant_states = \"True\"  #@param ['True', 'False']\n",
        "\n",
        "#@markdown draw with left click, hold shift for erasing\n",
        "\n",
        "if model_source != 'CHECKPOINT':\n",
        "  use_sample_pool, mutate_pool = {\n",
        "      '1 naive': (False, False),\n",
        "      '2 persistent': (True, False),\n",
        "      '3 mutating': (True, True)\n",
        "      }[model_type]\n",
        "  add_noise = add_noise == 'True'\n",
        "\n",
        "  model_str = get_model(\n",
        "      use_sample_pool=use_sample_pool, mutate_pool=mutate_pool,\n",
        "      loss_type=loss_type, add_noise=add_noise,\n",
        "      output='json')\n",
        "else:\n",
        "  last_checkpoint_fn = sorted(glob.glob('train_log/*.json'))[-1]\n",
        "  model_str = open(last_checkpoint_fn).read()\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "  window.SAMPLES = %s\n",
        "'''%(model_str, samples_str)\n",
        "\n",
        "display(IPython.display.Javascript(data_js))\n",
        "\n",
        "\n",
        "IPython.display.HTML('''\n",
        "<script src=\\\"https://unpkg.com/@tensorflow/tfjs@latest/dist/tf.min.js\\\"></script>\n",
        "<script src=\"https://cdnjs.cloudflare.com/ajax/libs/cash/4.1.2/cash.min.js\"></script>\n",
        "<div class=\"slidecontainer\">\n",
        "    brushSize:\n",
        "    <input type=\"range\" min=\"1\" max=\"10\" value=\"4\" class=\"slider\" id=\"brushSlider\">\n",
        "    <span id='radius'>2.5</span>\n",
        "</div>\n",
        "<canvas id='canvas' style=\"border: 1px solid black; image-rendering: pixelated;\"></canvas>\n",
        "<script>\n",
        "  \"use strict\";\n",
        "\n",
        "  // Adds the WASM backend to the global backend registry.\n",
        "  //import '@tensorflow/tfjs-backend-wasm';\n",
        "  // Set the backend to WASM and wait for the module to be ready.\n",
        "const main = async () => {\n",
        "\n",
        "  const sleep = (ms)=>new Promise(resolve => setTimeout(resolve, ms));\n",
        "  \n",
        "  const parseConsts = model_graph=>{\n",
        "    const dtypes = {'DT_INT32':['int32', 'intVal', Int32Array],\n",
        "                    'DT_FLOAT':['float32', 'floatVal', Float32Array]};\n",
        "    \n",
        "    const consts = {};\n",
        "    model_graph.modelTopology.node.filter(n=>n.op=='Const').forEach((node=>{\n",
        "      const v = node.attr.value.tensor;\n",
        "      const [dtype, field, arrayType] = dtypes[v.dtype];\n",
        "      if (!v.tensorShape.dim) {\n",
        "        consts[node.name] = [tf.scalar(v[field][0], dtype)];\n",
        "      } else {\n",
        "        const shape = v.tensorShape.dim.map(d=>parseInt(d.size));\n",
        "        let arr;\n",
        "        if (v.tensorContent) {\n",
        "          const data = atob(v.tensorContent);\n",
        "          const buf = new Uint8Array(data.length);\n",
        "          for (var i=0; i<data.length; ++i) {\n",
        "            buf[i] = data.charCodeAt(i);\n",
        "          }\n",
        "          arr = new arrayType(buf.buffer);\n",
        "        } else {\n",
        "          const size = shape.reduce((a, b)=>a*b);\n",
        "          arr = new arrayType(size);\n",
        "          arr.fill(v[field][0]);\n",
        "        }\n",
        "        consts[node.name] = [tf.tensor(arr, shape, dtype)];\n",
        "      }\n",
        "    }));\n",
        "    return consts;\n",
        "  }\n",
        "  \n",
        "  let paused = false;\n",
        "  let visibleChannel = -1;\n",
        "  let firingChance = 0.5;\n",
        "  let drawRadius = 2.5;\n",
        "\n",
        "  $('#brushSlider').on('input', e=>{\n",
        "      drawRadius = parseFloat(e.target.value)/2.0;\n",
        "      $('#radius').text(drawRadius);\n",
        "  });\n",
        "\n",
        "  const colorLookup = tf.tensor([\n",
        "      [128, 0, 0],\n",
        "      [230, 25, 75],\n",
        "      [70, 240, 240],\n",
        "      [210, 245, 60],\n",
        "      [250, 190, 190],\n",
        "      [170, 110, 40],\n",
        "      [170, 255, 195],\n",
        "      [165, 163, 159],\n",
        "      [0, 128, 128],\n",
        "      [128, 128, 0],\n",
        "      [0, 0, 0], // This is the default for digits.\n",
        "      [255, 255, 255] // This is the background.\n",
        "      ])\n",
        "\n",
        "  let backgroundWhite = true;\n",
        "\n",
        "\n",
        "  const run = async () => {\n",
        "      const r = await fetch(GRAPH_URL);\n",
        "      const consts = parseConsts(await r.json());\n",
        "\n",
        "      //const samples = tf.tensor(SAMPLES);\n",
        "      //console.log(samples);\n",
        "\n",
        "      const model = await tf.loadGraphModel(GRAPH_URL);\n",
        "\n",
        "      const samples = tf.tensor(SAMPLES);\n",
        "      console.log(samples.shape);\n",
        "      //const samples = tf.zeros([2,5, 28, 28]);\n",
        "\n",
        "      console.log(\"Loaded model\")\n",
        "      Object.assign(model.weights, consts);\n",
        "      // samples.gather(tf.range(0, 4, 1, 'int32')\n",
        "      const D = 28 * 2;\n",
        "      const state = tf.variable(tf.zeros([1, D, D, 20]));\n",
        "      const [_, h, w, ch] = state.shape;\n",
        "\n",
        "      const scale = 8;\n",
        "\n",
        "      const canvas = document.getElementById('canvas');\n",
        "      const ctx = canvas.getContext('2d');\n",
        "      canvas.width = w * scale;\n",
        "      canvas.height = h * scale;\n",
        "\n",
        "      const drawing_canvas = new OffscreenCanvas(w, h);\n",
        "      const draw_ctx = drawing_canvas.getContext('2d');\n",
        "\n",
        "      // Useful for understanding background color.\n",
        "      \n",
        "      //let blackAndWhite = tf.zeros();//.fill(0.01);\n",
        "      let arr = new Float32Array(h * w * 2);\n",
        "      arr.fill(0.01);\n",
        "      const blackAndWhiteFull = tf.tensor(arr, [1,h,w,2], tf.float32)\n",
        "\n",
        "      const drawCanvas = (imgd, e) => {\n",
        "          var matrix = [];\n",
        "          for(let i=0; i<imgd.width; i++) {\n",
        "              matrix[i] = [];\n",
        "              for(let j=0; j<imgd.height; j++) {\n",
        "                  let intensity = imgd.data[(imgd.height*j*4 + i*4)];\n",
        "                  // For drawing, we want to add shades of grey. For erasing, we don't.\n",
        "                  if (!e.shiftKey) {\n",
        "                    intensity *= (imgd.data[(imgd.height*j*4 + i*4 + 3)] / 255);\n",
        "                  }\n",
        "                  matrix[i][j] = intensity;\n",
        "              }\n",
        "          }\n",
        "\n",
        "          tf.tidy(() => {\n",
        "              const stroke = tf.tensor(matrix).transpose().toFloat().div(255.).expandDims(0).expandDims(3);\n",
        "              const stroke_pad = tf.concat([stroke, tf.zeros([1, h, w, ch-1])], 3);\n",
        "              const mask = tf.tensor(1.).sub(stroke);\n",
        "              if (e.shiftKey) {\n",
        "                  state.assign(state.mul(mask));\n",
        "              } else {\n",
        "                  state.assign(state.mul(mask).add(stroke_pad));\n",
        "              }\n",
        "          });\n",
        "\n",
        "          // Then clear the canvas.\n",
        "          draw_ctx.clearRect(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "      }\n",
        "\n",
        "      const line = (x0, y0, x1, y1, r, e) => {\n",
        "          draw_ctx.beginPath();\n",
        "          draw_ctx.moveTo(x0, y0);\n",
        "          draw_ctx.lineTo(x1, y1);\n",
        "          draw_ctx.strokeStyle = \"#ff0000\";\n",
        "          // Erasing has a much larger radius.\n",
        "          draw_ctx.lineWidth = (e.shiftKey ? 5. * r : r);\n",
        "          draw_ctx.stroke();\n",
        "\n",
        "          const imgd = draw_ctx.getImageData(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "          drawCanvas(imgd, e);\n",
        "      }\n",
        "\n",
        "\n",
        "      const circle = (x, y, r, e) => {\n",
        "          draw_ctx.beginPath();\n",
        "\n",
        "          const drawRadius = (e.shiftKey ? 5. * r : r) / 3.;\n",
        "\n",
        "          draw_ctx.arc(x, y, drawRadius, 0, 2 * Math.PI, false);\n",
        "          draw_ctx.fillStyle = \"#ff0000\";\n",
        "          draw_ctx.fill();\n",
        "          draw_ctx.lineWidth = 1;\n",
        "          draw_ctx.strokeStyle = \"#ff0000\";\n",
        "          draw_ctx.stroke();\n",
        "\n",
        "          const imgd = draw_ctx.getImageData(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "          drawCanvas(imgd, e);\n",
        "      }\n",
        "\n",
        "      const draw_r = 2.0;\n",
        "\n",
        "\n",
        "      const getClickPos = e=>{\n",
        "          const x = Math.floor((e.pageX-e.target.offsetLeft) / scale);\n",
        "          const y = Math.floor((e.pageY-e.target.offsetTop) / scale);\n",
        "          return [x, y];\n",
        "      }\n",
        "\n",
        "      let lastX = 0;\n",
        "      let lastY = 0;\n",
        "\n",
        "      canvas.onmousedown = e => {\n",
        "          const [x, y] = getClickPos(e);\n",
        "          lastX = x;\n",
        "          lastY = y;\n",
        "          circle(x,y,drawRadius, e);\n",
        "      }\n",
        "      canvas.onmousemove = e => {\n",
        "          const [x, y] = getClickPos(e);\n",
        "          if (e.buttons == 1) {\n",
        "              line(lastX,lastY, x,y,drawRadius, e);\n",
        "          }\n",
        "          lastX = x;\n",
        "          lastY = y;\n",
        "      }\n",
        "      const render = async () => {\n",
        "        if (!paused) {\n",
        "          tf.tidy(() => {\n",
        "              state.assign(model.execute(\n",
        "                { x: state,\n",
        "                  fire_rate: tf.tensor(firingChance),\n",
        "                  manual_noise: tf.randomNormal([1, h, w, ch-1], 0., 0.02)},\n",
        "                ['Identity']));\n",
        "          });\n",
        "        }\n",
        "        const imageData = tf.tidy(() => {\n",
        "            let rgbaBytes;\n",
        "            let rgba;\n",
        "            if (visibleChannel < 0) {\n",
        "                const isGray = state.slice([0,0,0,0],[1, h, w, 1]).greater(0.1).toFloat();\n",
        "                const isNotGray = tf.tensor(1.).sub(isGray);\n",
        "\n",
        "                const bnwOrder = backgroundWhite ?  [isGray, isNotGray] : [isNotGray, isGray];\n",
        "                let blackAndWhite = blackAndWhiteFull.mul(tf.concat(bnwOrder, 3));\n",
        "\n",
        "                const grey = state.gather([0], 3).mul(255);\n",
        "                const rgb = tf.gather(colorLookup,\n",
        "                                      tf.argMax(\n",
        "                                      tf.concat([\n",
        "                  state.slice([0,0,0,ch-10],[1,h,w,10]),\n",
        "                  blackAndWhite], 3), 3));\n",
        "\n",
        "                rgba = tf.concat([rgb, grey], 3)\n",
        "            } else {\n",
        "                rgba = state.gather([visibleChannel, visibleChannel, visibleChannel], 3)\n",
        "                  .pad([[0, 0], [0, 0], [0, 0], [0, 1]], 1).mul(255);\n",
        "            }\n",
        "            rgbaBytes = new Uint8ClampedArray(rgba.dataSync());\n",
        "\n",
        "            return new ImageData(rgbaBytes, w, h);\n",
        "        });\n",
        "        const image = await createImageBitmap(imageData);\n",
        "        //ctx.clearRect(0, 0, canvas.width, canvas.height);\n",
        "        ctx.fillStyle = backgroundWhite ? \"#ffffff\" : \"#000000\";\n",
        "        ctx.fillRect(0, 0, canvas.width, canvas.height);\n",
        "        ctx.imageSmoothingEnabled = false;\n",
        "        ctx.drawImage(image, 0, 0, canvas.width, canvas.height);\n",
        "\n",
        "        requestAnimationFrame(render);\n",
        "      }\n",
        "      render();\n",
        "  }\n",
        "\n",
        "  run();\n",
        "}\n",
        "main();\n",
        "  //tf.setBackend('wasm').then(() => main());\n",
        "\n",
        "  \n",
        "</script>\n",
        "''')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "HdniIRcvMtQo",
        "colab": {}
      },
      "source": [
        "#@title TensorFlow.js Demo Full whiteboard {run:\"auto\", vertical-output: true}\n",
        "#@markdown Select \"CHECKPOINT\" model to load the checkpoint created by running \n",
        "#@markdown cells from the \"Training\" section of this notebook.\n",
        "#@markdown Technical note: CE models should be rendered differently to avoid\n",
        "#@markdown black pixels showing for low magnitude.\n",
        "import IPython.display\n",
        "\n",
        "model_source = \"LOAD\"  #@param ['CHECKPOINT', 'LOAD']\n",
        "model_type = '3 mutating'  #@param ['1 naive', '2 persistent', '3 mutating']\n",
        "loss_type = \"l2\"  #@param ['l2', 'ce']\n",
        "add_noise = \"True\"  #@param ['True', 'False']\n",
        "#quant_states = \"True\"  #@param ['True', 'False']\n",
        "\n",
        "#@markdown draw with left click, hold shift for erasing\n",
        "\n",
        "if model_source != 'CHECKPOINT':\n",
        "  use_sample_pool, mutate_pool = {\n",
        "      '1 naive': (False, False),\n",
        "      '2 persistent': (True, False),\n",
        "      '3 mutating': (True, True)\n",
        "      }[model_type]\n",
        "  add_noise = add_noise == 'True'\n",
        "\n",
        "  model_str = get_model(\n",
        "      use_sample_pool=use_sample_pool, mutate_pool=mutate_pool,\n",
        "      loss_type=loss_type, add_noise=add_noise,\n",
        "      output='json')\n",
        "else:\n",
        "  last_checkpoint_fn = sorted(glob.glob('train_log/*.json'))[-1]\n",
        "  model_str = open(last_checkpoint_fn).read()\n",
        "\n",
        "data_js = '''\n",
        "  window.GRAPH_URL = URL.createObjectURL(new Blob([`%s`], {type: 'application/json'}));\n",
        "  window.SAMPLES = %s\n",
        "'''%(model_str, samples_str)\n",
        "\n",
        "display(IPython.display.Javascript(data_js))\n",
        "\n",
        "\n",
        "IPython.display.HTML('''\n",
        "<script src=\\\"https://unpkg.com/@tensorflow/tfjs@latest/dist/tf.min.js\\\"></script>\n",
        "<script src=\"https://cdnjs.cloudflare.com/ajax/libs/cash/4.1.2/cash.min.js\"></script>\n",
        "<div class=\"slidecontainer\">\n",
        "    brushSize:\n",
        "    <input type=\"range\" min=\"1\" max=\"10\" value=\"4\" class=\"slider\" id=\"brushSlider\">\n",
        "    <span id='radius'>2.5</span>\n",
        "</div>\n",
        "<canvas id='canvas' style=\"border: 1px solid black; image-rendering: pixelated;\"></canvas>\n",
        "<script>\n",
        "  \"use strict\";\n",
        "\n",
        "  // Adds the WASM backend to the global backend registry.\n",
        "  //import '@tensorflow/tfjs-backend-wasm';\n",
        "  // Set the backend to WASM and wait for the module to be ready.\n",
        "const main = async () => {\n",
        "\n",
        "  const sleep = (ms)=>new Promise(resolve => setTimeout(resolve, ms));\n",
        "  \n",
        "  const parseConsts = model_graph=>{\n",
        "    const dtypes = {'DT_INT32':['int32', 'intVal', Int32Array],\n",
        "                    'DT_FLOAT':['float32', 'floatVal', Float32Array]};\n",
        "    \n",
        "    const consts = {};\n",
        "    model_graph.modelTopology.node.filter(n=>n.op=='Const').forEach((node=>{\n",
        "      const v = node.attr.value.tensor;\n",
        "      const [dtype, field, arrayType] = dtypes[v.dtype];\n",
        "      if (!v.tensorShape.dim) {\n",
        "        consts[node.name] = [tf.scalar(v[field][0], dtype)];\n",
        "      } else {\n",
        "        const shape = v.tensorShape.dim.map(d=>parseInt(d.size));\n",
        "        let arr;\n",
        "        if (v.tensorContent) {\n",
        "          const data = atob(v.tensorContent);\n",
        "          const buf = new Uint8Array(data.length);\n",
        "          for (var i=0; i<data.length; ++i) {\n",
        "            buf[i] = data.charCodeAt(i);\n",
        "          }\n",
        "          arr = new arrayType(buf.buffer);\n",
        "        } else {\n",
        "          const size = shape.reduce((a, b)=>a*b);\n",
        "          arr = new arrayType(size);\n",
        "          arr.fill(v[field][0]);\n",
        "        }\n",
        "        consts[node.name] = [tf.tensor(arr, shape, dtype)];\n",
        "      }\n",
        "    }));\n",
        "    return consts;\n",
        "  }\n",
        "  \n",
        "  let paused = false;\n",
        "  let visibleChannel = -1;\n",
        "  let firingChance = 0.5;\n",
        "  let drawRadius = 2.5;\n",
        "\n",
        "  $('#brushSlider').on('input', e=>{\n",
        "      drawRadius = parseFloat(e.target.value)/2.0;\n",
        "      $('#radius').text(drawRadius);\n",
        "  });\n",
        "\n",
        "  const colorLookup = tf.tensor([\n",
        "      [128, 0, 0],\n",
        "      [230, 25, 75],\n",
        "      [70, 240, 240],\n",
        "      [210, 245, 60],\n",
        "      [250, 190, 190],\n",
        "      [170, 110, 40],\n",
        "      [170, 255, 195],\n",
        "      [165, 163, 159],\n",
        "      [0, 128, 128],\n",
        "      [128, 128, 0],\n",
        "      [0, 0, 0], // This is the default for digits.\n",
        "      [255, 255, 255] // This is the background.\n",
        "      ])\n",
        "\n",
        "  let backgroundWhite = true;\n",
        "\n",
        "\n",
        "  const run = async () => {\n",
        "      const r = await fetch(GRAPH_URL);\n",
        "      const consts = parseConsts(await r.json());\n",
        "\n",
        "      //const samples = tf.tensor(SAMPLES);\n",
        "      //console.log(samples);\n",
        "\n",
        "      const model = await tf.loadGraphModel(GRAPH_URL);\n",
        "\n",
        "      const samples = tf.tensor(SAMPLES);\n",
        "      //const samples = tf.zeros([2,5, 28, 28]);\n",
        "\n",
        "      console.log(\"Loaded model\")\n",
        "      Object.assign(model.weights, consts);\n",
        "      // samples.gather(tf.range(0, 4, 1, 'int32')\n",
        "      const D = 28 * 5;\n",
        "      const flatState = tf.concat([\n",
        "        samples.reshape([2,5, 28, 28])\n",
        "          .transpose([0,2,1,3]).reshape([1, 2 * 28, D]),\n",
        "        tf.zeros([1, D-(28*2), D])],1);\n",
        "      console.log(flatState)\n",
        "      const state = tf.variable(\n",
        "        tf.concat([flatState.expandDims(3),\n",
        "                    tf.zeros([1, D, D, 19])], 3));\n",
        "      const [_, h, w, ch] = state.shape;\n",
        "\n",
        "      const scale = 8;\n",
        "\n",
        "      const canvas = document.getElementById('canvas');\n",
        "      const ctx = canvas.getContext('2d');\n",
        "      canvas.width = w * scale;\n",
        "      canvas.height = h * scale;\n",
        "\n",
        "      const drawing_canvas = new OffscreenCanvas(w, h);\n",
        "      const draw_ctx = drawing_canvas.getContext('2d');\n",
        "\n",
        "      // Useful for understanding background color.\n",
        "      \n",
        "      //let blackAndWhite = tf.zeros();//.fill(0.01);\n",
        "      let arr = new Float32Array(h * w * 2);\n",
        "      arr.fill(0.01);\n",
        "      const blackAndWhiteFull = tf.tensor(arr, [1,h,w,2], tf.float32)\n",
        "\n",
        "      const drawCanvas = (imgd, e) => {\n",
        "          var matrix = [];\n",
        "          for(let i=0; i<imgd.width; i++) {\n",
        "              matrix[i] = [];\n",
        "              for(let j=0; j<imgd.height; j++) {\n",
        "                  let intensity = imgd.data[(imgd.height*j*4 + i*4)];\n",
        "                  // For drawing, we want to add shades of grey. For erasing, we don't.\n",
        "                  if (!e.shiftKey) {\n",
        "                    intensity *= (imgd.data[(imgd.height*j*4 + i*4 + 3)] / 255);\n",
        "                  }\n",
        "                  matrix[i][j] = intensity;\n",
        "              }\n",
        "          }\n",
        "\n",
        "          tf.tidy(() => {\n",
        "              const stroke = tf.tensor(matrix).transpose().toFloat().div(255.).expandDims(0).expandDims(3);\n",
        "              const stroke_pad = tf.concat([stroke, tf.zeros([1, h, w, ch-1])], 3);\n",
        "              const mask = tf.tensor(1.).sub(stroke);\n",
        "              if (e.shiftKey) {\n",
        "                  state.assign(state.mul(mask));\n",
        "              } else {\n",
        "                  state.assign(state.mul(mask).add(stroke_pad));\n",
        "              }\n",
        "          });\n",
        "\n",
        "          // Then clear the canvas.\n",
        "          draw_ctx.clearRect(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "      }\n",
        "\n",
        "      const line = (x0, y0, x1, y1, r, e) => {\n",
        "          draw_ctx.beginPath();\n",
        "          draw_ctx.moveTo(x0, y0);\n",
        "          draw_ctx.lineTo(x1, y1);\n",
        "          draw_ctx.strokeStyle = \"#ff0000\";\n",
        "          // Erasing has a much larger radius.\n",
        "          draw_ctx.lineWidth = (e.shiftKey ? 5. * r : r);\n",
        "          draw_ctx.stroke();\n",
        "\n",
        "          const imgd = draw_ctx.getImageData(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "          drawCanvas(imgd, e);\n",
        "      }\n",
        "\n",
        "\n",
        "      const circle = (x, y, r, e) => {\n",
        "          draw_ctx.beginPath();\n",
        "\n",
        "          const drawRadius = (e.shiftKey ? 5. * r : r) / 3.;\n",
        "\n",
        "          draw_ctx.arc(x, y, drawRadius, 0, 2 * Math.PI, false);\n",
        "          draw_ctx.fillStyle = \"#ff0000\";\n",
        "          draw_ctx.fill();\n",
        "          draw_ctx.lineWidth = 1;\n",
        "          draw_ctx.strokeStyle = \"#ff0000\";\n",
        "          draw_ctx.stroke();\n",
        "\n",
        "          const imgd = draw_ctx.getImageData(0, 0, draw_ctx.canvas.width, draw_ctx.canvas.height);\n",
        "          drawCanvas(imgd, e);\n",
        "      }\n",
        "\n",
        "      const draw_r = 2.0;\n",
        "\n",
        "\n",
        "      const getClickPos = e=>{\n",
        "          const x = Math.floor((e.pageX-e.target.offsetLeft) / scale);\n",
        "          const y = Math.floor((e.pageY-e.target.offsetTop) / scale);\n",
        "          return [x, y];\n",
        "      }\n",
        "\n",
        "      let lastX = 0;\n",
        "      let lastY = 0;\n",
        "\n",
        "      canvas.onmousedown = e => {\n",
        "          const [x, y] = getClickPos(e);\n",
        "          lastX = x;\n",
        "          lastY = y;\n",
        "          circle(x,y,drawRadius, e);\n",
        "      }\n",
        "      canvas.onmousemove = e => {\n",
        "          const [x, y] = getClickPos(e);\n",
        "          if (e.buttons == 1) {\n",
        "              line(lastX,lastY, x,y,drawRadius, e);\n",
        "          }\n",
        "          lastX = x;\n",
        "          lastY = y;\n",
        "      }\n",
        "      const render = async () => {\n",
        "        if (!paused) {\n",
        "          tf.tidy(() => {\n",
        "              state.assign(model.execute(\n",
        "                { x: state,\n",
        "                  fire_rate: tf.tensor(firingChance),\n",
        "                  manual_noise: tf.randomNormal([1, h, w, ch-1], 0., 0.02)},\n",
        "                ['Identity']));\n",
        "          });\n",
        "        }\n",
        "        const imageData = tf.tidy(() => {\n",
        "            let rgbaBytes;\n",
        "            let rgba;\n",
        "            if (visibleChannel < 0) {\n",
        "                const isGray = state.slice([0,0,0,0],[1, h, w, 1]).greater(0.1).toFloat();\n",
        "                const isNotGray = tf.tensor(1.).sub(isGray);\n",
        "\n",
        "                const bnwOrder = backgroundWhite ?  [isGray, isNotGray] : [isNotGray, isGray];\n",
        "                let blackAndWhite = blackAndWhiteFull.mul(tf.concat(bnwOrder, 3));\n",
        "\n",
        "                const grey = state.gather([0], 3).mul(255);\n",
        "                const rgb = tf.gather(colorLookup,\n",
        "                                      tf.argMax(\n",
        "                                      tf.concat([\n",
        "                  state.slice([0,0,0,ch-10],[1,h,w,10]),\n",
        "                  blackAndWhite], 3), 3));\n",
        "\n",
        "                rgba = tf.concat([rgb, grey], 3)\n",
        "            } else {\n",
        "                rgba = state.gather([visibleChannel, visibleChannel, visibleChannel], 3)\n",
        "                  .pad([[0, 0], [0, 0], [0, 0], [0, 1]], 1).mul(255);\n",
        "            }\n",
        "            rgbaBytes = new Uint8ClampedArray(rgba.dataSync());\n",
        "\n",
        "            return new ImageData(rgbaBytes, w, h);\n",
        "        });\n",
        "        const image = await createImageBitmap(imageData);\n",
        "        //ctx.clearRect(0, 0, canvas.width, canvas.height);\n",
        "        ctx.fillStyle = backgroundWhite ? \"#ffffff\" : \"#000000\";\n",
        "        ctx.fillRect(0, 0, canvas.width, canvas.height);\n",
        "        ctx.imageSmoothingEnabled = false;\n",
        "        ctx.drawImage(image, 0, 0, canvas.width, canvas.height);\n",
        "\n",
        "        requestAnimationFrame(render);\n",
        "      }\n",
        "      render();\n",
        "  }\n",
        "\n",
        "  run();\n",
        "}\n",
        "main();\n",
        "  //tf.setBackend('wasm').then(() => main());\n",
        "\n",
        "  \n",
        "</script>\n",
        "''')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}